{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import datasets\n",
    "import evaluate\n",
    "import json\n",
    "from datasets import load_dataset\n",
    "from cmrc_eval import evaluate_cmrc\n",
    "import collections\n",
    "from datasets import load_dataset, DatasetDict\n",
    "from transformers import DefaultDataCollator\n",
    "from transformers import AutoTokenizer, AutoModelForMultipleChoice, Trainer, TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MutipleChoicePredictor:\n",
    "    def __init__(self, model, tokenizer, device):\n",
    "        self.model = model\n",
    "        self.tokenizer = tokenizer\n",
    "        self.device = device\n",
    "    def pre_process(self,context,question,choices):\n",
    "        contexts,question_choice = [],[]\n",
    "        for choice in choices:\n",
    "            contexts.append(context)\n",
    "            question_choice.append(question + ' 选项列表:' + choice)\n",
    "                # 不足5个选项信息，补全到5个。\n",
    "        if len(choices) <5:\n",
    "            for _ in range(5-len(choices)):\n",
    "                contexts.append(context)\n",
    "                question_choice.append(question+\" 选项列表:\"+\"不知道\")\n",
    "        tokenized_datasets =  tokenizer(contexts,question_choice,truncation=\"only_first\",max_length=512,padding=\"max_length\",return_tensors=\"pt\")\n",
    "        return tokenized_datasets\n",
    "\n",
    "    def predict(self, inputs):\n",
    "        inputs =  {k:v.unsqueeze(0).to(self.device) for k,v in inputs.items()}\n",
    "        print(\"self.model(**inputs)self.model(**inputs)self.model(**inputs)self.model(**inputs)self.model(**inputs)\",inputs)\n",
    "        # print(self.model(**inputs))\n",
    "        return self.model(**inputs).logits\n",
    "\n",
    "    def post_process(self, logits,choices):\n",
    "        preidiction = torch.argmax(logits, dim=-1).cpu().item()\n",
    "        return choices[preidiction]\n",
    "\n",
    "    def __call__(self, context, question, choices):\n",
    "        inputs = self.pre_process(context, question, choices)\n",
    "        logits = self.predict(inputs)\n",
    "        result = self.post_process(logits, choices)\n",
    "        return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"/data/models/huggingface/chinese-macbert-large\")\n",
    "model  = AutoModelForMultipleChoice.from_pretrained(\"/data/models/huggingface/chinese-macbert-large\")\n",
    "pipe = MutipleChoicePredictor(model,tokenizer,\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe(\"小明在北京上班\",\"小明在哪里上班?\",[\"北京\",\"上海\",\"深圳\",\"广州\",\"成都\"])\n",
    "pipe(\"小明在广州上班\",\"小明在哪里上班?\",[\"北京\",\"上海\",\"深圳\",\"广州\",\"成都\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
